package heap.前K个高频元素;

import java.util.*;


public class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        int[] ans = new int[k];
        Map<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            map.put(nums[i], map.getOrDefault(nums[i], 0) + 1);
        }
        MinHeap<Integer[]> minHeap = new MinHeap<>((o1, o2) -> {
            if (o1[1] > o2[1]) {
                return 1;
            }
            if (o1[1] < o2[1]) {
                return -1;
            }
            return 0;
        });
        for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
            Integer number = entry.getKey();
            Integer count = entry.getValue();
            Integer[] temp = new Integer[]{number, count};
            minHeap.add(temp);
            if (minHeap.getSize() > k) {
                minHeap.removeHeapTop();
            }

        }
        for (int i = 0; i < k; i++) {
            ans[i] = minHeap.removeHeapTop()[0];
        }
        return ans;
    }

    /**
     * 小根堆
     * 底层是数组：
     * 包含：获取堆顶元素、堆顶元素出栈、节点上浮（新增）、节点下沉（堆顶出栈）、扩容
     */
    class MinHeap<T> {

        private List<T> heapArray;
        private int size;
        private Comparator<T> comparator;
        private int defaultCapacity = 10;

        public int getSize() {
            return this.size;
        }

        public MinHeap(Comparator<T> comparator) {
            // 修正：正确初始化
            this.heapArray = new ArrayList<>(defaultCapacity);
            this.size = 0;
            this.comparator = comparator;
        }

        public MinHeap(int capacity, Comparator<T> comparator) {
            this.heapArray = new ArrayList<>(capacity);
            this.size = 0;
            this.comparator = comparator;
        }

        /**
         * 获取父节点
         *
         * @param index
         * @return
         */
        private int parent(int index) {
            return (index - 1) / 2;
        }

        /**
         * 获取左节点
         *
         * @param index
         * @return
         */
        private int left(int index) {
            return index * 2 + 1;
        }

        /**
         * 获取右节点
         *
         * @param index
         * @return
         */
        private int right(int index) {
            return index * 2 + 2;
        }


        /**
         * 节点上浮
         *
         * @param index
         */
        private void floatUp(int index) {
            while (index > 0 && comparator.compare(heapArray.get(index), heapArray.get(parent(index))) == -1) {
                swap(index, parent(index));
                index = parent(index);
            }
        }

        private void swap(int i, int j) {
            T temp = heapArray.get(i);
            heapArray.set(i, heapArray.get(j));
            heapArray.set(j, temp);
        }

        /**
         * 节点下沉
         */
        private void sink(int index) {
            int minIndex = index;
            int left = left(index);
            int right = right(index);

            // 找出当前节点、左孩子、右孩子三者中的最小值
            if (left < size && comparator.compare(heapArray.get(left), heapArray.get(minIndex)) == -1) {
                minIndex = left;
            }
            if (right < size && comparator.compare(heapArray.get(right), heapArray.get(minIndex)) == -1) {
                minIndex = right;
            }
            // 如果最小值不是当前节点，说明需要下沉
            if (minIndex != index) {
                swap(index, minIndex);
                sink(minIndex); // 递归地继续下沉
            }
        }

        public void add(T number) {
            heapArray.add(number);
            floatUp(size);
            size++;
        }

        /**
         * 删除栈顶元素
         *
         * @return
         */
        public T removeHeapTop() {
            if (size == 0) {
                throw new IllegalStateException("Heap is empty");
            }
            T min = heapArray.get(0);
            heapArray.set(0, heapArray.get(--size));
            // 修正：移除最后一个元素（现在是重复的）
            if (size < heapArray.size()) {
                heapArray.remove(size);
            }
            sink(0);
            return min;
        }

        public T peek() {
            return heapArray.get(0);
        }

    }

    public static void main(String[] args) {
        Solution solution = new Solution();
        System.out.println(solution.topKFrequent(new int[]{1,1,1,2,2,3} , 2 ));
    }
}
